{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Hyperelasticity "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![hyperelasticity.png](figures/hyperelasticity.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "using JuAFEM\n",
    "using Tensors\n",
    "using KrylovMethods\n",
    "using TimerOutputs\n",
    "import ProgressMeter\n",
    "const ∇ = Tensors.gradient;"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### NeoHook Material"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "immutable NeoHook{T}\n",
    "    μ::T\n",
    "    λ::T\n",
    "end\n",
    "\n",
    "function compute_2nd_PK(mp::NeoHook, E)\n",
    "    I = one(E)\n",
    "    C = 2E + one(E)\n",
    "    invC = inv(C)\n",
    "    J = sqrt(det(C))\n",
    "    return mp.μ *(I - invC) + mp.λ * log(J) * invC\n",
    "end\n",
    "\n",
    "function constitutive_driver(mp::NeoHook, E)\n",
    "    ∂S∂E, SPK = ∇(E -> compute_2nd_PK(mp, E), E, :all)\n",
    "    return SPK, ∂S∂E\n",
    "end;"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Assembler routines"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Loop over all cells \n",
    "function assemble{dim}(grid::Grid{dim}, dh::DofHandler, K, f, cv, fv, mp, u)\n",
    "    n = ndofs_per_cell(dh)\n",
    "    Ke = zeros(n, n)\n",
    "    fe = zeros(n)\n",
    "\n",
    "    assembler = start_assemble(K, f)\n",
    "\n",
    "    # loop over all cells in the grid\n",
    "    @timeit \"assemble\" for cell in CellIterator(dh)\n",
    "        # reset\n",
    "        fill!(Ke, 0)\n",
    "        fill!(fe, 0)\n",
    "\n",
    "        global_dofs = celldofs(cell)\n",
    "        ue = u[global_dofs] # element dofs\n",
    "        @timeit \"inner assemble\" assemble_element!(Ke, fe, cell, cv, fv, mp, ue)\n",
    "\n",
    "        assemble!(assembler, global_dofs, fe, Ke)\n",
    "    end\n",
    "\n",
    "    return f, K\n",
    "end;"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Assembles the contribution from the cell to ke and fe\n",
    "function assemble_element!(ke, fe, cell, cv, fv, mp, ue)\n",
    "    b = Vec{3}((0.0, -0.5, 0.0))\n",
    "    t = Vec{3}((0.1, 0.0, 0.0))\n",
    "    ndofs = getnbasefunctions(cv)\n",
    "    reinit!(cv, cell)\n",
    "    fill!(ke, 0.0)\n",
    "    fill!(fe, 0.0)\n",
    "    δE = Vector{SymmetricTensor{2, 3, eltype(ue), 6}}(ndofs)\n",
    "\n",
    "    for qp in 1:getnquadpoints(cv)\n",
    "        ∇u = function_gradient(cv, qp, ue)\n",
    "        dΩ = getdetJdV(cv, qp)\n",
    "\n",
    "        # strain and stress + tangent\n",
    "        F = one(∇u) + ∇u\n",
    "        E = symmetric(1/2 * (F' ⋅ F - one(F)))\n",
    "\n",
    "        S, ∂S∂E = constitutive_driver(mp, E)\n",
    "\n",
    "        # Hoist computations of δE\n",
    "        for i in 1:ndofs\n",
    "            δFi = shape_gradient(cv, qp, i)\n",
    "            δE[i] = symmetric(1/2*(δFi'⋅F + F'⋅δFi))\n",
    "        end\n",
    "\n",
    "        for i in 1:ndofs\n",
    "            δFi = shape_gradient(cv, qp, i)\n",
    "            δu = shape_value(cv, qp, i)\n",
    "            fe[i] += (δE[i] ⊡ S) * dΩ\n",
    "            fe[i] -= (δu ⋅ b) * dΩ\n",
    "            δE∂S∂E = δE[i] ⊡ ∂S∂E\n",
    "            S∇δu = S ⋅ δFi'\n",
    "            for j in 1:ndofs\n",
    "                δ∇uj = shape_gradient(cv, qp, j)\n",
    "                ke[i, j] += (δE∂S∂E ⊡ δE[j] + S∇δu ⊡ δ∇uj' ) * dΩ\n",
    "            end\n",
    "        end\n",
    "    end\n",
    "\n",
    "    for face in 1:nfaces(cell)\n",
    "        if onboundary(cell, face)\n",
    "            reinit!(fv, cell, face)\n",
    "            for q_point in 1:getnquadpoints(fv)\n",
    "                dΓ = getdetJdV(fv, q_point)\n",
    "                for i in 1:ndofs\n",
    "                    δu = shape_value(fv, q_point, i)\n",
    "                    fe[i] -= (δu ⋅ t) * dΓ\n",
    "                end\n",
    "            end\n",
    "        end\n",
    "    end\n",
    "end;"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Main solver routine"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "solve (generic function with 1 method)"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "function solve()\n",
    "    reset_timer!()\n",
    "\n",
    "    const dim = 3\n",
    "\n",
    "    # Generate a grid\n",
    "    N = 10\n",
    "    L = 1.0\n",
    "    left = zero(Vec{dim})\n",
    "    right = L * ones(Vec{dim})\n",
    "    grid = generate_grid(Tetrahedron, ntuple(x->N, dim), left, right)\n",
    "\n",
    "    # Node sets\n",
    "    addnodeset!(grid, \"clamped\", x -> norm(x[1]) ≈ 1)\n",
    "    addnodeset!(grid, \"rotation\", x -> norm(x[1]) ≈ 0)\n",
    "\n",
    "    # Material parameters\n",
    "    E = 10.0\n",
    "    ν = 0.3\n",
    "    μ = E / (2(1 + ν))\n",
    "    λ = (E * ν) / ((1 + ν) * (1 - 2ν))\n",
    "    mp = NeoHook(μ, λ)\n",
    "\n",
    "    # finite element base\n",
    "    ip = Lagrange{dim, RefTetrahedron, 1}()\n",
    "    qr = QuadratureRule{dim, RefTetrahedron}(1)\n",
    "    qr_face = QuadratureRule{dim-1, RefTetrahedron}(1)\n",
    "    cv = CellVectorValues(qr, ip)\n",
    "    fv = FaceVectorValues(qr_face, ip)\n",
    "\n",
    "    # DofHandler\n",
    "    dh = DofHandler(grid)\n",
    "    push!(dh, :u, dim) # Add a displacement field\n",
    "    close!(dh)\n",
    "\n",
    "    function rotation(X, t, θ = deg2rad(60.0))\n",
    "        x, y, z = X\n",
    "        return t * Vec{dim}(\n",
    "            (0.0,\n",
    "            L/2 - y + (y-L/2)*cos(θ) - (z-L/2)*sin(θ),\n",
    "            L/2 - z + (y-L/2)*sin(θ) + (z-L/2)*cos(θ)\n",
    "            ))\n",
    "    end\n",
    "\n",
    "    dbc = DirichletBoundaryConditions(dh)\n",
    "    # Add a homogenoush boundary condition on the \"clamped\" edge\n",
    "    add!(dbc, :u, getnodeset(grid, \"clamped\"), (x,t) -> [0.0, 0.0, 0.0], collect(1:dim))\n",
    "    add!(dbc, :u, getnodeset(grid, \"rotation\"), (x,t) -> rotation(x, t), collect(1:dim))\n",
    "    close!(dbc)\n",
    "    t = 0.5\n",
    "    update!(dbc, t)\n",
    "\n",
    "    println(\"Analysis with \", length(grid.cells), \" elements\")\n",
    "\n",
    "    # pre-allocate\n",
    "    _ndofs = ndofs(dh)\n",
    "    un = zeros(_ndofs) # previous solution vector\n",
    "    u  = zeros(_ndofs)\n",
    "    Δu = zeros(_ndofs)\n",
    "\n",
    "    apply!(un, dbc)\n",
    "\n",
    "    K = create_sparsity_pattern(dh)\n",
    "    f = zeros(_ndofs)\n",
    "\n",
    "    newton_itr = -1\n",
    "    NEWTON_TOL = 1e-8\n",
    "    prog = ProgressMeter.ProgressThresh(NEWTON_TOL, \"Solving:\")\n",
    "\n",
    "    while true; newton_itr += 1\n",
    "        u .= un .+ Δu\n",
    "        f, K = assemble(grid, dh, K, f, cv, fv, mp, u)\n",
    "        normg = norm(f[JuAFEM.free_dofs(dbc)])\n",
    "        apply_zero!(K, f, dbc)\n",
    "        ProgressMeter.update!(prog, normg; showvalues = [(:iter, newton_itr)])\n",
    "\n",
    "        if normg < NEWTON_TOL\n",
    "            break\n",
    "        end\n",
    "\n",
    "        if newton_itr > 30\n",
    "            error(\"Reached maximum Newton iterations, aborting\")\n",
    "            break\n",
    "        end\n",
    "\n",
    "        @timeit \"linear solve\" ΔΔu, flag, relres, iter, resvec = cg(K, f; maxIter = 1000, tol = min(1e-3, normg))\n",
    "        @assert flag == 0\n",
    "\n",
    "        apply_zero!(ΔΔu, dbc)\n",
    "        Δu .-= ΔΔu\n",
    "    end\n",
    "\n",
    "    # save the solution\n",
    "    @timeit \"export\" begin\n",
    "        vtkfile = vtk_grid(\"hyperelasticity\", dh)\n",
    "        vtk_point_data(vtkfile, dh, u)\n",
    "        vtk_save(vtkfile)\n",
    "    end\n",
    "\n",
    "    print_timer(linechars = :ascii)\n",
    "    return u\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Analysis with 5000 elements\n",
      " \u001b[1m---------------------------------------------------------------------------\u001b[22m\n",
      "                                    Time                   Allocations      \n",
      "                            ----------------------   -----------------------\n",
      "      Tot / % measured:          203ms / 78.2%           63.2MiB / 45.2%    \n",
      "\n",
      " Section            ncalls     time   %tot     avg     alloc   %tot      avg\n",
      " ---------------------------------------------------------------------------\n",
      " assemble                6   99.9ms  63.0%  16.7ms   26.1MiB  91.4%  4.35MiB\n",
      "   inner assemble    30.0k   62.9ms  39.6%  2.10μs   19.2MiB  67.3%     672B\n",
      " linear solve            5   51.2ms  32.3%  10.2ms    683KiB  2.34%   137KiB\n",
      " export                  1   7.52ms  4.74%  7.52ms   1.79MiB  6.26%  1.79MiB\n",
      " \u001b[1m---------------------------------------------------------------------------\u001b[22m"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\r",
      "Solving: (thresh = 1e-08, value = 0.000384834)\u001b[34m\n",
      "  iter:  3\u001b[39m\u001b[1G\u001b[K\u001b[A\r",
      "Solving: Time: 0:00:00 (6 iterations)\u001b[34m\n",
      "  iter:  5\u001b[39m\n"
     ]
    }
   ],
   "source": [
    "u = solve();"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Hyperelasticity passed!\n"
     ]
    }
   ],
   "source": [
    "Base.Test.@test norm(u) ≈ 4.870833706518008\n",
    "println(\"Hyperelasticity passed!\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "name": "julia-0.6",
   "display_name": "Julia 0.6.0",
   "language": "julia"
  },
  "language_info": {
   "mimetype": "application/julia",
   "file_extension": ".jl",
   "version": "0.6.0",
   "name": "julia"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
